#include "kernel_operator.h"
#include "op_common.h"
#include <cstdint>

using namespace AscendC;

template <typename _DT_X, typename _DT_SIZE, typename _DT_STRIDE,
          typename _DT_STORAGE_OFFSET, typename _DT_Y>
class Kernel {
public:
  TPipe pipe;
  ShapeData xShape;
  ShapeData sizeStrideShape;
  ShapeData yShape;
  int64_t y_shape_data[64] = {0};
  int64_t stride_data[64] = {0};
  int64_t storageOffset = 0;
  DefGlobalTensor(X);
  DefGlobalTensor(SIZE);
  DefGlobalTensor(STRIDE);
  DefGlobalTensor(STORAGE_OFFSET);
  DefGlobalTensor(Y);

  DefBufVECIN(INDEX);
  DefBufVECCALC(STRIDE);
  DefBufVECCALC(WORK_LOCAL);
  DefBufVECCALC(TEMP);
  DefOutQue(DATAINDEX);

public:
  __aicore__ inline Kernel() {}

  template <typename T>
  __aicore__ inline void Init(GM_ADDR x, GM_ADDR size, GM_ADDR stride,
                              GM_ADDR storage_offset, GM_ADDR y,
                              T tiling_data) {
    xShape.init(tiling_data.xShape);
    sizeStrideShape.init(tiling_data.sizeStrideShape);
    GSetBuffer(X, x, 0, xShape.size);
    GSetBuffer(SIZE, size, 0, sizeStrideShape.size);
    GSetBuffer(STRIDE, stride, 0, sizeStrideShape.size);
    GSetBuffer(STORAGE_OFFSET, storage_offset, 0, 1);

    // y_shape_data[0] = sizeStrideShape.size;
    // y_shape_data[1] = 0;
    // for (auto i = 0; i < sizeStrideShape.size; ++i) {
    //   y_shape_data[i + 2] = GTensorName(SIZE).GetValue(i);
    //   stride_data[i] = GTensorName(STRIDE).GetValue(i);
    // }
    // yShape.init(y_shape_data);
    GSetBuffer(Y, y, 0, 1024 *4);
    // storageOffset = GTensorName(STORAGE_OFFSET).GetValue(0);

    // auto calcLength = alignToBlock<float>(yShape.dimNum);
    // auto bufLength = calcLength * sizeof(float);
    InitTBufBuffer(INDEX, 1024 *4);
    // InitTBufBuffer(STRIDE, bufLength);
    // InitTBufBuffer(TEMP, bufLength);
    // InitTBufBuffer(WORK_LOCAL, workLocalSizeBtyte<float>(bufLength));
    // InitQueue(DATAINDEX, bufLength);

    // TBufGet(STRIDE, float);
    // for (auto i = 0; i < yShape.dimNum; ++i) {
    //   BTensor(STRIDE)(i) = stride_data[i];
    // }
    // TBufGet(INDEX, float);
    // MY_DUP_INT32(BTensor(INDEX), zeroFloatInt, yShape.dimNum);
  }

  __aicore__ inline void Process() {
    for (auto i = 0; i < yShape.size; ++i) {
      int64_t indexs[32] = {0};
      yShape.Offset2Indexs(i, indexs);
      auto xOffset = storageOffset;
      for (auto j = 0; j < yShape.dimNum; ++j) {
        xOffset += indexs[j] * stride_data[j];
      }
      GTensorName(Y).SetValue(i, GTensorName(X).GetValue(xOffset));
    }
  }

  __aicore__ inline void Process2() {
    TBufGet(INDEX, float);
    TBufGet(STRIDE, float);
    TBufGet(WORK_LOCAL, float);
    TBufGet(TEMP, float);
    for (auto i = 0; i < yShape.size; ++i) {
      {
        QueAlloc(DATAINDEX, float);
        // Mul(BTensor(TEMP), BTensor(INDEX), BTensor(STRIDE), yShape.dimNum);
        // ReduceSum(LTensor(DATAINDEX), BTensor(TEMP), BTensor(WORK_LOCAL),
        //           yShape.dimNum);
        // for (auto j = yShape.dimNum - 1; j >= 0; --j) {
        //   auto newIndex = BTensor(INDEX).GetValue(j) + 1.0f;
        //   auto maxIndex = yShape[j];
        //   if (newIndex < maxIndex) {
        //     BTensor(INDEX).SetValue(j, newIndex);
        //     break;
        //   }
        //   newIndex = int32_t(newIndex) % maxIndex;
        //   BTensor(INDEX).SetValue(j, float(newIndex));
        // }
        EnQue(DATAINDEX);
      }
      {
        DeQue(DATAINDEX, float);
        // auto index = LTensor(DATAINDEX).GetValue(0) + storageOffset;
        // GTensorName(Y).SetValue(i, GTensor(X).GetValue(int32_t(index)));
        QueFree(DATAINDEX);
      }
    }
  }

  __aicore__ inline void Process3() {
    TBufGet(INDEX, _DT_X);
    for (auto i = 0; i < 60000; ++i) {
      DataCopy(GTensor(Y), BTensor(INDEX), 1024 *4 /sizeof(_DT_X));
    }
  }
};

extern "C" __global__ __aicore__ void
as_strided(GM_ADDR x, GM_ADDR size, GM_ADDR stride, GM_ADDR storage_offset,
           GM_ADDR y, GM_ADDR workspace, GM_ADDR tiling) {
  GET_TILING_DATA(tiling_data, tiling);
  Kernel<DTYPE_X, DTYPE_SIZE, DTYPE_STRIDE, DTYPE_STORAGE_OFFSET, DTYPE_Y> op;
  op.Init(x, size, stride, storage_offset, y, tiling_data);
  op.Process3();
}